package beast.evolution.tree;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;

import beast.core.Description;
import beast.core.Input;
import beast.evolution.alignment.Alignment2;
import beast.evolution.alignment.Sequence;
import beast.evolution.datatype.Nucleotide;

@Description("This class extends Tree and allows nodes to store sequence info")
public class Tree2 extends Tree {
	public Input<Alignment2> m_pData = new Input<Alignment2>("aligment2",
			"same alignment data that is used for tree construction");

	protected int sequenceSize;

	/**
	 * default access modifier, seen only to itself
	 */
	Node2[] m_nodes;

	protected beast.evolution.tree.Node2 root2;

    public Tree2() {}
	
	/**
	 * This constructor takes "tree" and "alignment" as input use root2 (type
	 * "Node2") as the new root for Tree2 create nodes for Tree2 by traverse the
	 * Tree and replace Node with Node2
	 */
	public Tree2(Tree tree, Alignment2 alignment2) throws Exception {
		/** assign alignment **/
		m_pData.setValue(alignment2, this);
		/** copy node count info to Tree2 **/
		nodeCount = tree.nodeCount;
		internalNodeCount = tree.internalNodeCount;
		leafNodeCount = tree.leafNodeCount;
		/** create "root2" from "root" in "Tree" **/
		root2 = new Node2(tree.root);
	}

	/**
	 * Assign initial values of sequence info to Leaf nodes based on Alignment2
	 * Assign zero values to internal nodes and root
	 */
	public void initSequences() {
		/** get the length of the sequences **/
		sequenceSize = m_pData.get().getSiteCount();
		/** convert tree to array expression **/
		m_nodes = new Node2[nodeCount];
		listNodes(root2, m_nodes);
		/**
		 * for all nodes in the array, assign default value to the sequence
		 * parameter assign observed sequences to leaf nodes based on Taxon
		 * match
		 */
		for (int i = 0; i < m_nodes.length; i++) {
			if (!m_nodes[i].isLeaf()) {
				m_nodes[i].setSeq(new ArrayList<Integer>(Collections.nCopies(
						sequenceSize, 0)));
			} else {
				try {
					m_nodes[i].setSeq(m_pData.get().getSequenceByID(
							m_nodes[i].getID()));
				} catch (Exception e1) {
					// to do
				}
				System.out.println(m_nodes[i].getID());
				/**
				 * a void function that takes Node2 and Alignment2 as input and
				 * assigns real sequence value to Node2's integer array to be
				 * done something like this:
				 * m_nodes[i].setSeq(getSequenceByID(m_nodes[i].getID()));
				 * getSequenceByID() is a getter function in Alignment2
				 */
			}

		}

		for (int i = 0; i < m_nodes.length; i++) {
			Nucleotide nucleo = new Nucleotide();
			List<Integer> tmp_seq = m_nodes[i].getSeq();
			String real_seq = nucleo.state2string(tmp_seq);
			System.out.println(real_seq);
		}

	}

	void listNodes(Node2 node, Node2[] nodes) {
		nodes[node.getNr()] = node;
		node.m_tree = this;
		if (!node.isLeaf()) {
			listNodes(node.m_left, nodes);
			if (node.m_right != null) {
				listNodes(node.m_right, nodes);
			}
		}
	}

	public void showSequences() throws Exception {
		Nucleotide nucleo = new Nucleotide();
		for (int i = 0;i<m_nodes.length;i++) {
			List<Integer> tmp_seq = m_nodes[i].getSeq();
			/**
			 * System.out.println("The length of the sequence is:" +
			 * aligned_seq.getSiteCount());
			 */
			String real_seq = nucleo.state2string(tmp_seq);
			System.out.println("The node is:" + i + " and the Seq is: \n"+ real_seq);
		}

	}

	public void PupkoAllSites() {
		// apply PupkoOneSite to all sites
		for(int i = 0;i<sequenceSize;i++) {
			this.PupkoOneSite(i);
		}
	}
	
	// the Pupko algorithm
	public void PupkoOneSite(int siteNum) {

		/* check whether NucleotideSampling is working correctly */
		// double [] nucleoFreq = new double [] {0.1,0.2,0.3,0.4};
		// for (int i=0;i<999;i++) {
		// System.out.println(this.NucleotideSampling(nucleoFreq));
		// }

		// create a converter
		// Nucleotide nucleo = new Nucleotide();

		// get alignment from m_pData
		// Alignment2 aligned_seq = m_pData.get();

		// temp: site number
		//int siteNum = 7;

		/*
		 * temp: an array to temp store new char state Left to right: A C G T =>
		 * 0 1 2 3
		 */
		int[] newCharState = new int[] { 1, 1, 1, 1 };

		// create an array to store nucleotide states at the temp site for each
		// node
		int[] m_states = new int[m_nodes.length];
		for (int i = 0; i < m_nodes.length; i++) {
			m_states[i] = m_nodes[i].getSeq().get(siteNum);
			// System.out.println("The node is:" + m_nodes[i].getNr()
			// + " and the nucleotide state is" + m_states[i]);
		}

		/* deal with leaf nodes first (since they can be readily calculated) */
		for (int i = 0; i < root2.getLeafNodeCount(); i++) {
			/*
			 * given the observed nucleo type j at leaf node, set the jth column
			 * of the pMatrix
			 */
			this.setLeafpMatrix(m_states[i], m_nodes[i].getpMatrix(),
					m_nodes[i].getLength());
			/*
			 * given the observed nucleo type at leaf node i, we determine
			 * "newCharState" for this leaf node
			 */
			this.setLeafCharState(m_states[i], newCharState);
			/* update the CharState of this leaf node */
			m_nodes[i].setCharState(newCharState);
			System.out.println("The species name is:"
					+ m_nodes[i].getID()
					+ "\n The node is:"
					+ m_nodes[i].getNr()
					+ " and the nucleotide state is"
					+ Arrays.toString(newCharState)
					+ "\n and the branch length is:"
					+ m_nodes[i].getLength()
					+ " and the pMatrix is: \n"
					+ Arrays.deepToString(m_nodes[i].getpMatrix()).replace(
							"], ", "\n").replaceAll(",|\\[|\\]", ""));
		}

		// create a list to store numbers for internal nodes
		List<Integer> internalNodeNum = new ArrayList<Integer>();
		for (int i = root2.getLeafNodeCount(); i < root2.getNodeCount(); i++) {
			internalNodeNum.add(i);
		}
		// remove the root number from internalNodeNum
		internalNodeNum.remove(internalNodeNum.indexOf(root2.getNr()));

		/*
		 * while internalNodeNum is not empty, constantly search through
		 * internal node list and find those nodes whose children have already
		 * been calculated
		 */

		while (internalNodeNum.size() > 0) {
			for (int i = 0; i < internalNodeNum.size(); i++) {
				if (m_nodes[internalNodeNum.get(i)].getLength() > 0) {
					// get the node number of the left child node and right
					// child
					// node
					int leftNum = m_nodes[internalNodeNum.get(i)].m_left
							.getNr();
					int rightNum = m_nodes[internalNodeNum.get(i)].m_right
							.getNr();
					// if both left and right child node has been updated, then
					// update this internal node
					if ((m_nodes[leftNum].getCharState()[0] != -1)
							&& (m_nodes[rightNum].getCharState()[0] != -1)) {
						this
								.setInternalpMatrix(m_nodes[internalNodeNum
										.get(i)]);
						System.out.println("remove element:"
								+ internalNodeNum.get(i)
								+ "and its left child is:"
								+ m_nodes[leftNum].getNr()
								+ "and its right child is:"
								+ m_nodes[rightNum].getNr());
						// it is important to update the state of this node
						// under
						// consideration, otherwise, you will get error
						this.setInternalCharState(m_nodes[internalNodeNum.get(i)]);
						internalNodeNum.remove(i);
					}
				}
				if (m_nodes[internalNodeNum.get(i)].getLength() == 0) {
					// get the node number of the left child node and right
					// child
					// node
					int leftNum = m_nodes[internalNodeNum.get(i)].m_left
							.getNr();
					int rightNum = m_nodes[internalNodeNum.get(i)].m_right
							.getNr();
					// if both left and right child node has been updated, then
					// update this internal node
					if ((m_nodes[leftNum].getCharState()[0] != -1)
							&& (m_nodes[rightNum].getCharState()[0] != -1)) {
						System.out.println("remove element:"
								+ internalNodeNum.get(i)
								+ "and its left child is:" + leftNum
								+ "and its right child is:" + rightNum
								+ "and its extra child is:"
								+ root2.m_left.getNr());
						// it is important to update the state of this node
						// under
						// consideration, otherwise, you will get error
						this.setSudoRootCharState(m_nodes[internalNodeNum.get(i)]);
						internalNodeNum.remove(i);
					}
				}
			}
		}
		
		// trace back and set CharState for internal nodes
		this.traceBackCharStates(siteNum, m_nodes[root2.m_right.getNr()]);
		
		// initialize CharStates for all nodes
		for (int i=0;i<m_nodes.length;i++) {
			m_nodes[i].setCharStateNegative();
		}
		
	}// end of Pupko

	/* this function is used to set Char state for leaf nodes */
	public void setLeafCharState(int LeafChar, int[] LeafCharState) {
		for (int i = 0; i < LeafCharState.length; i++) {
			LeafCharState[i] = LeafChar;
		}
	}
	
	/* set Char State for internal nodes */
	public void setInternalCharState(Node2 internalNode) {
		int [] newCharState = new int [] {-1,-1,-1,-1};
		/* this array is used to store temp probability vector of the child CharState j (A,T,C,G) given the parent CharState i*/
		double [] pChildStateGivenParent = new double [4];
		// loop through all ancestral states
		for (int i=0; i < 4; i++) {
			// calculate pChildStateGivenParent
			for (int j=0; j < 4; j++) {
				pChildStateGivenParent[j] =  internalNode.getpMatrix()[i][j] / internalNode.getpMatrixRowSum(i);
			}
			// sample the newCharState for this internal node based on pChildStateGivenParent
			newCharState[i] = this.NucleotideSampling(pChildStateGivenParent);
			System.out.println("The p vector is: "+ Arrays.toString(pChildStateGivenParent) + " and the nucleo sampled is: " + newCharState[i]);
		}
		// set CharState for this internal node finally
		internalNode.setCharState(newCharState);
	}

	/* set root char state */
	public void setSudoRootCharState(Node2 sudoRootNode) {
		// int to store the sudoRootState
		int sudoRoot = -1;
		
		// get the node number of the left child node and right
		// child
		// node
		int leftNum = sudoRootNode.m_left.getNr();
		int rightNum = sudoRootNode.m_right.getNr();
		int outgroupNum = root2.m_left.getNr();
		
		double [] pChildStateGivenParent = new double [4];
		/* vector to store nucleotide frequency 
		 * Here, since we are using JC69, frequencies will cancel out
		 * but if we use other model and calculate frequencies based on our data,
		 * we may set this value at the beginning*/
		double [] NucleoFreq = new double [] {0.25,0.25,0.25,0.25};
		double sudoRootProbSum = 0.0;
		for (int j=0;j<4;j++) {
			//System.out.println("j is:"+j+" and prob is: " + NucleoFreq[j] * m_nodes[leftNum].getpMatrixRowSum(j) * m_nodes[rightNum].getpMatrixRowSum(j) * m_nodes[outgroupNum].getpMatrixRowSum(j));
			sudoRootProbSum += NucleoFreq[j] * m_nodes[leftNum].getpMatrixRowSum(j) * m_nodes[rightNum].getpMatrixRowSum(j) * m_nodes[outgroupNum].getpMatrixRowSum(j);
		}
		for (int j =0;j<4;j++) {
			pChildStateGivenParent[j] = (NucleoFreq[j] * m_nodes[leftNum].getpMatrixRowSum(j) * m_nodes[rightNum].getpMatrixRowSum(j) * m_nodes[outgroupNum].getpMatrixRowSum(j))/sudoRootProbSum;
		}
		sudoRoot = this.NucleotideSampling(pChildStateGivenParent);
		sudoRootNode.setCharState(new int [] {sudoRoot, sudoRoot,sudoRoot,sudoRoot});
		System.out.println("The p vector is: "+ Arrays.toString(pChildStateGivenParent) + " and the root nucleo sampled is: " + sudoRoot);
	}
	
	/* trace back from the root to leafs and set the Char States for all nodes at the given site */
	public void traceBackCharStates(int siteNum, Node2 internalNode) {
		// deal with sudoRootNode
		if(!internalNode.isLeaf()) {
			int parentNum = internalNode.m_Parent.getNr();
			int leftNum = internalNode.m_left.getNr();
			int rightNum = internalNode.m_right.getNr();
			
			if(internalNode.m_Parent.isRoot()) {
				internalNode.setSeqSite(siteNum, internalNode.charState[0]);
				System.out.println("This sudoRoot Num is: " + internalNode.getNr()  + " and state is: "+internalNode.getSeqSite(siteNum));
				this.traceBackCharStates(siteNum, m_nodes[leftNum]);
				this.traceBackCharStates(siteNum, m_nodes[rightNum]);	
			}else {
				// deal with other internal nodes
				internalNode.setSeqSite(siteNum, internalNode.charState[m_nodes[parentNum].getSeqSite(siteNum)]);
				System.out.println("This internal Node is: " + internalNode.getNr() + " and its state is: "+ internalNode.getSeqSite(siteNum) );
				this.traceBackCharStates(siteNum, m_nodes[leftNum]);
				this.traceBackCharStates(siteNum, m_nodes[rightNum]);	
			}
		}else {
			System.out.println("Reach leaf: "+ internalNode.getNr());
		}
	}
	
	/* this func is used to set pMatrix for leaf nodes */
	public void setLeafpMatrix(int LeafChar, double[][] pMatrix,
			double branchLength) {
		for (int i = 0; i < 4; i++) {
			pMatrix[i][LeafChar] = this.transitProbJC69(i, LeafChar,
					branchLength);
		}
	}

	/* set pMatrix for internal nodes */
	public void setInternalpMatrix(Node2 internalNode) {
		double P_ij = 0.0;
		int leftNum = internalNode.m_left.getNr();
		int rightNum = internalNode.m_right.getNr();
		for (int i = 0; i < 4; i++) {
			for (int j = 0; j < 4; j++) {
				P_ij = this.transitProbJC69(i, j, internalNode.getLength());
				P_ij = P_ij * m_nodes[leftNum].getpMatrixRowSum(j)
						* m_nodes[rightNum].getpMatrixRowSum(j);
				internalNode.setpMatrix(i, j, P_ij);
			}
		}
		System.out.println("The internal node is:"
				+ internalNode.getNr()
				+ "\n and the branch length is:"
				+ internalNode.getLength()
				+ " and the pMatrix is: \n"
				+ Arrays.deepToString(internalNode.getpMatrix()).replace("], ",
						"\n").replaceAll(",|\\[|\\]", ""));
	}

	/*
	 * This function might be temporary: it calculate the probability that
	 * nucleotide i at the beginning of the branch and nucleotide j at the end
	 * of the branch based on JC69 model
	 */
	public double transitProbJC69(int i, int j, double branchLength) {
		double transitProb = 0.0;
		if (i == j) {
			transitProb = 0.25 * (1 + 3 * Math.exp((-4.0 / 3.0) * branchLength));
		} else {
			transitProb = 0.25 * (1 - Math.exp((-4.0 / 3.0) * branchLength));
		}
		return transitProb;
	}

	/*
	 * sample nucleotide based on discrete probability distribution this
	 * function may form its own class and move to somewhere else This function
	 * implements inverse sampling in the discrete case. For example:
	 */
	// X p(X) F(X)
	// 0 0.1 0.1
	// 1 0.3 0.4=0.1 + 0.3
	// 2 0.2 0.6=0.4+0.2
	// 3 0.4 1=0.6+0.4
	/*
	 * The first outcome for which the random deviate is smaller than (or is
	 * equal to) the associated cumulative probability corresponds to the
	 * sampled outcome
	 */

	public int NucleotideSampling(double[] pNucleotide) {
		// create an integer to store final result
		int sampledNucleo = -1;
		// from p(X) we calculate the CDF F(X)
		double[] fNucleotide = new double[4];
		fNucleotide[0] = pNucleotide[0];
		for (int i = 1; i < pNucleotide.length; i++) {
			fNucleotide[i] = fNucleotide[i - 1] + pNucleotide[i];
		}
		// draw a random number between 0 and 1
		double randNum = Math.random();
		// get the first nucleotide whose F(X) is larger or equal to randNum
		for (int i = 0; i < fNucleotide.length; i++) {
			if (fNucleotide[i] > randNum) {
				sampledNucleo = i;
				break;
			}
		}
		return sampledNucleo;
	}
	
} // class Tree2